import cv2
import argparse
from utils import load_model, draw_boxes


def process_video(video_path, model):
    """
    处理视频文件，使用模型对视频帧进行对象检测，并显示结果。
    :param video_path: 视频文件的路径
    :param model: 加载的YOLO模型
    """
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"Error opening video file: {video_path}")
        return

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        results = model(frame)  # 对每一帧进行预测
        frame = draw_boxes(results, frame)  # 在帧上绘制边界框和标签

        cv2.imshow('Video Detection', frame)
        if cv2.waitKey(1) == ord('q'):  # 按 'q' 键退出
            break

    cap.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="YOLOv9 视频对象检测")
    parser.add_argument('--video', help='视频文件路径', default='data/input_video.mp4')
    args = parser.parse_args()

    model = load_model('models/yolov9_model.pt')  # 加载模型
    process_video(args.video, model)